#!/usr/bin/python

# extract_questions.py

# .pot format:
#
#     <whitespace>
#     #: file:line ...
#     msgid "unstranslated_string"
#     msgstr "translated_string"

# We need:
#
#   category, q_name (with file:line references)
#     choices (with file:line references)


import sys
import os
import collections
import functools
import re

Debug = True

from doctest_tools import setpath
Root_dir = setpath.setpath(__file__, remove_first=True)[0]
#print "Root_dir", Root_dir

from naimath.engine import extracted_questions
from naimath.lib import po_headers

Translation_dir = os.path.join(Root_dir, 'translation', '0.1', 'po')
#print "Translation_dir", Translation_dir

def read_file(filename, d_results, q_results):
    r'''Update d_results, q_results dict with entries from filename.

    D_results dict is {untranslated_diagnosis: [(file, line)...]}

    Q_results dict is {category: {q_name:
                                     ([(file, line)...],
                                      {choice: [(file, line)...]}
                                     ),
                                 },
                      }
    '''
    if Debug: print "reading", filename
    with open(filename) as f:
        line_cont = None
        for lineno, line in enumerate(f, start=1):
            line = line.strip()
            if line_cont is not None:
                line_cont = \
                  line_cont(line, filename, lineno, d_results, q_results)
            else:
                #print "checking %s(%d): %s" % (filename, lineno, line)
                for pattern, fn in Patterns:
                    m = pattern.match(line)
                    if m:
                        #print "got match"
                        if Debug: print "matched", line
                        line_cont = fn(m, filename, lineno,
                                       d_results, q_results)

def make_q_results():
    r'''Makes a result dict with the proper default entries.

        >>> d = make_q_results()
        >>> d['cat']['a'][1]['b']
        []
        >>> d['cat']['a']
        ([], defaultdict(<type 'list'>, {'b': []}))
    '''
    return collections.defaultdict(
             lambda: collections.defaultdict(
                       lambda: ([], collections.defaultdict(list))))

def q_a(ans_fn, m, filename, lineno, d_results, q_results):
    rest = m.group(1).strip()
    if rest:
        return get_question(ans_fn, rest, filename, lineno,
                            d_results, q_results)
    return functools.partial(get_questions, ans_fn)

Question_re = re.compile(
        r'''\s* ['"] ( [^"']+ ) ["'] \s* ,
            \s* ['"] ( [^"']+ ) ['"] \s* ,
            (.*)''',
        re.VERBOSE)

def get_question(get_ans_fn, line, filename, lineno, d_results, q_results):
    m = Question_re.match(line)
    assert m, "%s(%d): invalid question: %s" % (filename, lineno, line)
    category = m.group(1).strip()
    question = m.group(2).strip()
    rest = m.group(3).strip()
    q_results[category][question][0].append((filename, lineno))
    if rest:
        return get_ans_fn(category, question, rest, filename, lineno,
                          d_results, q_results)
    return functools.partial(get_ans_fn, category, question)

Single_ans_re = re.compile(
        r'''\s* ['"] ( [^'"]+ ) ['"] \s* ,''',
        re.VERBOSE)

def single_ans(category, question, ans, filename, lineno, d_results, q_results):
    m = Single_ans_re.match(ans)
    assert m, "%s(%d): invalid answer: %s" % (filename, lineno, ans)
    choice = m.group(1).strip()
    q_results[category][question][1][choice].append((filename, lineno))

Ans_tuple_re = re.compile(
        r'''\s* ( [^= ]+ ) \s* = \s* (.*)''',
        re.VERBOSE)

def ans_tuple(category, question, ans, filename, lineno, d_results, q_results):
    m = Ans_tuple_re.match(ans)
    assert m, "%s(%d): invalid answer: %s" % (filename, lineno, ans)
    choice = m.group(1).strip()
    q_results[category][question][1][choice].append((filename, lineno))
    rest = m.group(2).strip()
    if rest:
        return skip_score(category, question, rest, filename, lineno,
                          d_results, q_results)
    return functools.partial(skip_score, category, question)

Score_re = re.compile(r'''[-+\d.]+ (.*)''', re.VERBOSE)

def skip_score(category, question, score, filename, lineno, d_results,
               q_results):
    m = Score_re.match(score)
    assert m, "%s(%d): invalid score: %s" % (filename, lineno, score)
    rest = m.group(1).strip()
    if rest[0] == ')': return None
    rest = rest[1:].lstrip()
    if rest:
        return ans_tuple(category, question, rest, filename, lineno,
                         d_results, q_results)
    return functools.partial(ans_tuple, category, question)

def diagnosis(m, filename, lineno, d_results, q_results):
    name = m.group(1)
    d_results[name].append((filename, lineno))

Patterns = (
    (re.compile(r'use\s+diagnosis\((\w+),'), diagnosis),
    (re.compile(r'score\.got\((.*)'), functools.partial(q_a, ans_tuple)),
    (re.compile(r'score\.not_\((.*)'), functools.partial(q_a, single_ans)),
)

def dump_d(d_results):
    for d, refs in d_results.iteritems():
        print "%s: %s" % (d, ' '.join("%s:%d" % (filename, lineno)
                                      for filename, lineno in refs))

def dump_q(q_results):
    for category, questions in q_results.iteritems():
        print "%s:" % (category,)
        for q, (refs, choices) in questions.iteritems():
            print "%s: %s" % (q, ' '.join("%s:%d" % (filename, lineno)
                                          for filename, lineno in refs))
            for choice, refs in choices.iteritems():
                print "    %s: %s" % (choice,
                                      ' '.join("%s:%d" % (filename, lineno)
                                               for filename, lineno in refs))

def write_d_pot(d_results, filename):
    with open(os.path.join(Translation_dir, filename), 'w') as f:
        po_headers.write_header(f)
        for d in sorted(d_results.iterkeys()):
            print >> f
            refs = d_results[d]
            refs.sort()
            print >> f, "#:", ' '.join("%s:%d" %
                                         (os.path.basename(file), lineno)
                                       for file, lineno in refs)
            print >> f, 'msgid "%s"' % (d,)
            print >> f, 'msgstr "%s"' % (d,)

def write_extracted_questions(q_results):
    if Debug: print "write_extracted_questions"
    all_q_names = set()
    for category, questions in q_results.iteritems():
        for q_name, (_, choices) in questions.iteritems():
            q_info = find_q_name(q_name, category, guess_q_type(choices.keys()))
            all_choices = set()
            for choice, references in choices.iteritems():
                choice_info = find_choice(q_info, choice)
                refs = ' '.join("%s:%d" % (os.path.basename(filename), lineno)
                                for filename, lineno in sorted(references))
                if len(choice_info) < 2: choice_info.append(refs)
                else: choice_info[1] = refs
                all_choices.add(choice)
            for c in q_info[3:]:
                if c[0] not in all_choices and len(c) > 1:
                    c[1] = ''
            all_q_names.add((q_name, category))
    to_delete = []
    for i, q in enumerate(extracted_questions.Questions):
        if (q[0], q[1]) not in all_q_names: to_delete.append(i)
    for i in to_delete[::-1]:
        del extracted_questions.Questions[i]
    rewrite_extracted_questions()

def guess_q_type(choices):
    choices = frozenset(choices)
    if choices.issubset(('yes', 'no')): return "yes_no"
    if choices.issubset(('positive', 'negative')): return "positive_negative"
    return "single_choice"

def find_q_name(q_name, category, q_type):
    for q in extracted_questions.Questions:
        if q_name == q[0] and category == q[1]: return q
    ans = [q_name, category, q_type]
    extracted_questions.Questions.append(ans)
    return ans

def find_choice(q_info, choice):
    for c in q_info[3:]:
        if c[0] == choice: return c
    ans = [choice]
    q_info.append(ans)
    return ans

def rewrite_extracted_questions():
    if Debug: print "rewrite_extracted_questions"
    with open(extracted_questions.__file__, 'w') as f:
        print >> f, "# extracted_questions.py"
        print >> f
        print >> f, "Questions = ["
        for q_info in sorted(extracted_questions.Questions):
            print >> f, "    [%r, %r, %r," % tuple(q_info[:3])
            for choice in q_info[3:]:
                print >> f, "        %r," % (choice,)
            print >> f, "      ],"
        print >> f, "]"

def usage():
    sys.stderr.write("usage: extract_questions.py krb_file...\n")
    sys.exit(2)

if __name__ == "__main__":
    if len(sys.argv) < 2: usage()
    d_results = collections.defaultdict(list)
    q_results = make_q_results()
    for filename in sys.argv[1:]:
        read_file(filename, d_results, q_results)
    #dump_d(d_results)
    #dump_q(q_results)
    write_d_pot(d_results, "diagnosis.pot")
    write_extracted_questions(q_results)
